import requests
import os
import base64
from openai import OpenAI
from bs4 import BeautifulSoup
from PIL import Image, ImageDraw

# ========== 1. 图片下载与保存 ==========
img_url = "https://ofasys-multimodal-wlcb-3-toshanghai.oss-cn-shanghai.aliyuncs.com/Qwen3VL/demo/omni_parsing/179729.jpg"
response = requests.get(img_url)
img_name = os.path.basename(img_url)
with open(img_name, "wb") as f:
    f.write(response.content)

# ========== 2. 图片Base64编码 ==========
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

base64_image = encode_image(img_name)

# ========== 3. 调用通义千问API（OpenAI兼容模式） ==========
client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

prompt = """请生成带精确边界框的HTML，每个元素标注“data-bbox='x1,y1,x2,y2'”（逗号分隔），
严格区分标题、文本、表格、图表，输出格式与示例一致：
<div class="h1" data-bbox="100,50,800,80">宇通客车 (600066)</div>
"""

completion = client.chat.completions.create(
    model="qwen-vl-max",
    messages=[
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
                {"type": "text", "text": prompt}
            ],
        }
    ],
    response_format={"type": "text"},
    timeout=120,
)

output = completion.choices[0].message.content
print("带边界框的原始HTML结果：\n", output)

# ========== 4. 边界框可视化 ==========
def draw_bbox_html(image_path, html_content, output_path="解析结果_带边界框.jpg"):
    img = Image.open(image_path)
    draw = ImageDraw.Draw(img)
    soup = BeautifulSoup(html_content, "html.parser")
    bbox_elements = soup.find_all(attrs={"data-bbox": True})

    if not bbox_elements:
        print("⚠️  未检测到边界框坐标，直接保存原图")
        img.save(output_path)
        return

    for elem in bbox_elements:
        bbox_str = elem.get("data-bbox", "").strip()
        if not bbox_str:
            continue
        try:
            x1, y1, x2, y2 = map(int, bbox_str.split(','))
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(img.width, x2), min(img.height, y2)
            draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
        except ValueError:
            print(f"⚠️  无效坐标格式：{bbox_str}，已跳过")

    img.save(output_path)
    print(f"✅ 边界框可视化图已保存至：{os.path.abspath(output_path)}")

draw_bbox_html(img_name, output)